# -*- coding: utf-8 -*-
"""
Created on Tue Oct 24 09:44:25 2017
"""

import xml.etree.ElementTree as ET
import os
import shutil

'''计算文件夹下文件个数'''


def fileCountIn(dir):
    return sum([len(files) for root, dirs, files in os.walk(dir)])


def copy_from_testsets(save_dir, test_file_list):
    copied_img_1 = 0
    copied_img_2 = 0
    copied_img_3 = 0
    copied_img_4 = 0

    with open(test_file_list) as f:
        for line in f.readlines():
            line = line.strip()
            xml_dir = line.split('.')[0] + '.xml'
            xml_file = xml_dir.replace("JPEGImages", "Annotations")
            tree = ET.parse(xml_file)
            root = tree.getroot()
            for obj in root.iter('object'):
                if obj.find('name').text == 'person' and copied_img_1 < 500:
                    xml_name = os.path.split(xml_file)[1]
                    image_name = os.path.split(line)[1]
                    image_path = line
                    save_annotation_dir = os.path.join(save_dir + 'pedestrian', "Annotations")
                    save_image_dir = os.path.join(save_dir + 'pedestrian', "JPEGImages")
                    # save_yolov3_labels = os.path.join(save_dir, "labels")

                    if not os.path.exists(save_annotation_dir):
                        os.makedirs(save_annotation_dir)
                    if not os.path.exists(save_image_dir):
                        os.makedirs(save_image_dir)

                    dst_xml_full_path = os.path.join(save_annotation_dir, xml_name)
                    dst_img_full_path = os.path.join(save_image_dir, image_name)
                    shutil.copyfile(xml_file, dst_xml_full_path)
                    shutil.copyfile(image_path, dst_img_full_path)
                    copied_img_1 += 1
                    if copied_img_1 % 100 == 0:
                        print("pedstration: {}".format(copied_img_1))
                    break
                elif obj.find('name').text == 'vehicle' and copied_img_2 < 500:
                    xml_name = os.path.split(xml_file)[1]
                    image_name = os.path.split(line)[1]
                    image_path = line

                    save_annotation_dir = os.path.join(save_dir + obj.find('name').text, "Annotations")
                    save_image_dir = os.path.join(save_dir + obj.find('name').text, "JPEGImages")
                    # save_yolov3_labels = os.path.join(save_dir, "labels")

                    if not os.path.exists(save_annotation_dir):
                        os.makedirs(save_annotation_dir)
                    if not os.path.exists(save_image_dir):
                        os.makedirs(save_image_dir)

                    dst_xml_full_path = os.path.join(save_annotation_dir, xml_name)
                    dst_img_full_path = os.path.join(save_image_dir, image_name)
                    shutil.copyfile(xml_file, dst_xml_full_path)
                    shutil.copyfile(image_path, dst_img_full_path)
                    copied_img_2 += 1
                    if copied_img_2 % 100 == 0:
                        print("vehicle: {}".format(copied_img_2))
                    break
                elif obj.find('name').text == 'ship' and copied_img_3 < 1000:
                    xml_name = os.path.split(xml_file)[1]
                    image_name = os.path.split(line)[1]
                    image_path = line

                    save_annotation_dir = os.path.join(save_dir + obj.find('name').text, "Annotations")
                    save_image_dir = os.path.join(save_dir + obj.find('name').text, "JPEGImages")
                    # save_yolov3_labels = os.path.join(save_dir, "labels")

                    if not os.path.exists(save_annotation_dir):
                        os.makedirs(save_annotation_dir)
                    if not os.path.exists(save_image_dir):
                        os.makedirs(save_image_dir)

                    dst_xml_full_path = os.path.join(save_annotation_dir, xml_name)
                    dst_img_full_path = os.path.join(save_image_dir, image_name)
                    shutil.copyfile(xml_file, dst_xml_full_path)
                    shutil.copyfile(image_path, dst_img_full_path)
                    copied_img_3 += 1
                    if copied_img_3 % 100 == 0:
                        print("ship: {}".format(copied_img_2))
                    break
                elif obj.find('name').text == 'fighter' and copied_img_4 < 1000:
                    xml_name = os.path.split(xml_file)[1]
                    image_name = os.path.split(line)[1]
                    image_path = line

                    save_annotation_dir = os.path.join(save_dir + "plane", "Annotations")
                    save_image_dir = os.path.join(save_dir + "plane", "JPEGImages")
                    # save_yolov3_labels = os.path.join(save_dir, "labels")

                    if not os.path.exists(save_annotation_dir):
                        os.makedirs(save_annotation_dir)
                    if not os.path.exists(save_image_dir):
                        os.makedirs(save_image_dir)

                    dst_xml_full_path = os.path.join(save_annotation_dir, xml_name)
                    dst_img_full_path = os.path.join(save_image_dir, image_name)
                    shutil.copyfile(xml_file, dst_xml_full_path)
                    shutil.copyfile(image_path, dst_img_full_path)
                    copied_img_4 += 1
                    if copied_img_4 % 100 == 0:
                        print("plane: {}".format(copied_img_2))
                    break
            if copied_img_1 >= 499 and copied_img_2 >= 499 and copied_img_3 >=999 and copied_img_4 >= 999:
                break


def main():
    xml_dir = '/home/ccj/dataset/vehicle_person/COCO/voc/Annotations'
    save_dir = '/home/ccj/dataset/test-sets/'

    test_file_list = '/home/dataset/zhangyan/15_yolov3_4/test_img_path.txt'
    copy_from_testsets(save_dir, test_file_list)
    return 0

    xml_Count = fileCountIn(xml_dir)
    print('xml Count in dir:', xml_Count)
    chg_num = 0
    copied_img_1 = 0
    copied_img_2 = 0
    object_dicts = {}
    for root, dirs, files in os.walk(xml_dir):
        for file in files:
            xml_file = os.path.join(xml_dir, file)
            # change_xml_node(xml_file)
            # objects = find_objects(xml_file)

            tree = ET.parse(xml_file)
            root = tree.getroot()
            for obj in root.iter('object'):
                if obj.find('name').text == 'person' and copied_img_1 < 500:
                    xml_name = os.path.split(xml_file)[1]
                    image_name = xml_name.split('.')[0] + '.jpg'
                    image_path = os.path.join(os.path.split(xml_file)[0].replace('Annotations', 'JPEGImages'),
                                              image_name)
                    save_annotation_dir = os.path.join(save_dir+'pedestrian', "Annotations")
                    save_image_dir = os.path.join(save_dir+'pedestrian', "JPEGImages")
                    # save_yolov3_labels = os.path.join(save_dir, "labels")

                    if not os.path.exists(save_annotation_dir):
                        os.makedirs(save_annotation_dir)
                    if not os.path.exists(save_image_dir):
                        os.makedirs(save_image_dir)

                    dst_xml_full_path = os.path.join(save_annotation_dir, xml_name)
                    dst_img_full_path = os.path.join(save_image_dir, image_name)
                    shutil.copyfile(xml_file, dst_xml_full_path)
                    shutil.copyfile(image_path, dst_img_full_path)
                    copied_img_1 += 1
                    print("pedstration: {}".format(copied_img_1))
                    break
                elif obj.find('name').text == 'vehicle' and copied_img_2 < 500:
                    xml_name = os.path.split(xml_file)[1]
                    image_name = xml_name.split('.')[0] + '.jpg'
                    image_path = os.path.join(os.path.split(xml_file)[0].replace('Annotations', 'JPEGImages'),
                                              image_name)

                    save_annotation_dir = os.path.join(save_dir+obj.find('name').text, "Annotations")
                    save_image_dir = os.path.join(save_dir+obj.find('name').text, "JPEGImages")
                    # save_yolov3_labels = os.path.join(save_dir, "labels")

                    if not os.path.exists(save_annotation_dir):
                        os.makedirs(save_annotation_dir)
                    if not os.path.exists(save_image_dir):
                        os.makedirs(save_image_dir)

                    dst_xml_full_path = os.path.join(save_annotation_dir, xml_name)
                    dst_img_full_path = os.path.join(save_image_dir, image_name)
                    shutil.copyfile(xml_file, dst_xml_full_path)
                    shutil.copyfile(image_path, dst_img_full_path)
                    copied_img_2 += 1
                    print("vehicle: {}".format(copied_img_2))
                    break
        if copied_img_1 >= 499 and copied_img_2 >= 499:
            break
            # chg_num = chg_num + 1
            # if chg_num % 1000 == 0:
            #     print('completed ', chg_num * 100 / xml_Count, '% ...')
    print('completed ', chg_num * 100 / xml_Count, '% ...')
    print('total:', chg_num)
    print('copied person image: ', copied_img_1)
    print('copied vehicle image: ', copied_img_2)


if __name__ == '__main__':
    main()
